import rasterio
# from rasterio.transform import Affine
import numpy as np

def read_tif(file_path):
    with rasterio.open(file_path) as src:
        data = src.read(1).astype(np.float32)  # 将数据转换为 float32 类型
        meta = src.meta
        # 将无效值转换为 nan
        nodata_value = src.nodata
        if nodata_value is not None:
            data[data == nodata_value] = np.nan
        return data, meta

def calculate_K(SAN, SIL, CLA, C, SNI):
    # 确保所有输入都是有效的（非 nan）
    valid_mask = ~np.isnan(SAN) & ~np.isnan(SIL) & ~np.isnan(CLA) & ~np.isnan(C) & ~np.isnan(SNI)

    K = np.full_like(SAN, np.nan, dtype=np.float32)
    K[valid_mask] = (
        (0.2 + 0.3 * np.exp(0.0256 * SAN[valid_mask] * (1 - SIL[valid_mask] / 100))) *
        ((SIL[valid_mask] / (CLA[valid_mask] + SIL[valid_mask])) ** 0.3) *
        (1 - (0.25 * C[valid_mask]) / (C[valid_mask] + np.exp(3.72 - 2.59 * C[valid_mask]))) *
        (1 - (0.7 * SNI[valid_mask]) / (SNI[valid_mask] + np.exp(-5.51 + 22.95 * SNI[valid_mask]))) *
        0.1317
    )
    return K

def save_as_tif(data, meta, output_path):
    # Update the metadata to reflect the number of bands and data type
    meta.update(
        driver='GTiff',
        count=1,
        dtype=data.dtype,
        nodata= -9999  # 设置无效值为 nan
    )

    with rasterio.open(output_path, 'w', **meta) as dst:
        dst.write(data, 1)

# 读取tif文件中的变量数据及其元数据
path_to_SAN = r'F:\code\dev\gep-calculation-helper\SF\input\SR_SF_inpu\T_sand_china.tif'
path_to_SIL = r'F:\code\dev\gep-calculation-helper\SF\input\SR_SF_inpu\T_silt_china.tif'
path_to_CLA = r'F:\code\dev\gep-calculation-helper\SF\input\SR_SF_inpu\T_clay_china.tif'
path_to_C = r'F:\code\dev\gep-calculation-helper\SF\input\SR_SF_inpu\ORG_CARBON1_china_1km.tif'

SAN, meta_SAN = read_tif(path_to_SAN)
SIL, meta_SIL = read_tif(path_to_SIL)
CLA, meta_CLA = read_tif(path_to_CLA)
C, meta_C = read_tif(path_to_C)
SNI = 1 - SAN/100

# 确保所有变量具有相同的形状和元数据
assert SAN.shape == SIL.shape == CLA.shape == C.shape == SNI.shape, "All input rasters must have the same dimensions"
meta = meta_SAN  # 使用其中一个变量的元数据作为输出文件的元数据

# 计算K值
K = calculate_K(SAN, SIL, CLA, C, SNI)

# 将结果保存为tif文件
output_path = './output/SR_K_china_1km.tif'
save_as_tif(K, meta, output_path)

print(f"K value saved to {output_path}")
